package ai.people.core.grpc.utils;

import ai.people.core.grpc.lib.error.Error;
import ai.people.netmon.framework.exception.type.BusinessException;
import com.google.protobuf.Message;
import io.grpc.Metadata;
import io.grpc.Status;
import io.grpc.protobuf.ProtoUtils;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.function.Function;
import java.util.function.Supplier;

/**
 * @author yuanqinglong
 * @date 2022/6/5 13:05
 */
public class StreamObserverDelegate<ReqT extends Message, RespT extends Message> implements StreamObserver<RespT> {
    private static final Logger logger = LoggerFactory.getLogger(StreamObserverDelegate.class);
    private final StreamObserver<RespT> originResponseObserver;

    public StreamObserverDelegate(StreamObserver<RespT> originResponseObserver) {
        this.originResponseObserver = originResponseObserver;
    }

    public static <ReqT extends Message, RespT extends Message> StreamObserverDelegate<ReqT, RespT> build(StreamObserver<RespT> originResponseObserver) {
        return new StreamObserverDelegate<>(originResponseObserver);
    }

    @Override
    public void onNext(RespT value) {
        this.originResponseObserver.onNext(value);
    }

    @Override
    public void onError(Throwable t) {
        logger.error("grpc request error:", t);
        if (t instanceof BusinessException) {
            // 业务异常，返回错误码和默认文案到客户端
            BusinessException serviceException = (BusinessException) t;
            Metadata trailers = new Metadata();
            Error errorInfo = Error.newBuilder()
                    .setCode(serviceException.getCode())
                    .setMessage(serviceException.getMessage())
                    .build();
            Metadata.Key<Error> ERROR_INFO_TRAILER_KEY =
                    ProtoUtils.keyForProto(errorInfo);
            trailers.put(ERROR_INFO_TRAILER_KEY, errorInfo);
            this.originResponseObserver.onError(Status.FAILED_PRECONDITION
                    .withCause(serviceException)
                    .asRuntimeException(trailers));
        } else {
            // 非业务异常，返回异常详情到客户端。
            this.originResponseObserver.onError(Status.UNKNOWN
                    // 这里就是自定义异常信息
                    .withDescription(t.getMessage())
                    .withCause(t)
                    .asRuntimeException());
        }
        // 抛出异常让上层业务感知(比如事务回滚等可能要用到)
        throw new RuntimeException(t);
    }

    @Override
    public void onCompleted() {
        if (originResponseObserver != null) {
            originResponseObserver.onCompleted();
        }
    }

    /**
     * 执行业务(自动处理异常)
     *
     * @author masaiqi
     * @date 2021/4/12 18:11
     */
    public RespT execute(Function<ReqT, RespT> function, ReqT request) {
        RespT response = null;
        try {
            response = function.apply(request);
        } catch (Throwable e) {
            this.onError(e);
        } finally {
            this.onNext(response);
            this.onCompleted();
        }
        return response;
    }

    /**
     * 执行业务(自动处理异常)
     *
     * @author masaiqi
     * @date 2021/4/12 18:11
     */
    public RespT execute(Supplier<RespT> supplier) {
        RespT response = null;
        try {
            response = supplier.get();
        } catch (Throwable e) {
            this.onError(e);
        } finally {
            this.onNext(response);
            this.onCompleted();
        }
        return response;
    }
}
